from typing import Tuple
from PIL import Image
import torch
from torch import Tensor
import torch.nn.functional as F
from torchvision.transforms import ToTensor, ToPILImage

from compressai.ops import compute_padding

import numpy as np

""" image Tensor Operation """
def load_image(filepath: str) -> Image.Image:
    return Image.open(filepath).convert("RGB")


def img2torch(img: Image.Image) -> torch.Tensor:
    return ToTensor()(img).unsqueeze(0)


def torch2img(x: torch.Tensor) -> Image.Image:
    return ToPILImage()(x.clamp_(0, 1).squeeze())


def pad(x, p=2**6, *kargs, **kwargs):
    h, w = x.size(-2), x.size(-1)
    pad, _ = compute_padding(h, w, min_div=p)
    return F.pad(x, pad, *kargs)


def crop(x, size, *kargs, **kwargs):
    H, W = x.size(-2), x.size(-1)
    h, w = size
    _, unpad = compute_padding(h, w, out_h=H, out_w=W)
    return F.pad(x, unpad, *kargs)
                 

